# -*- coding:utf-8 -*-
"""
Created on 2011-2-8

@author: summit
"""
import vtk
import vtkJolly
import struct

executing = 0

def execute_limited(pipeline, updateExtent=None, wantWholeDataset=None,
                    memLimit=512, noLimits=False, alwaysSplit=False, 
                    numberOfDivisions=1):
    global executing
    
    if noLimits or (not memLimit and not alwaysSplit):
        streamer = vtk.vtkImageDataStreamer()
        streamer.SetNumberOfStreamDivisions(1)
    
    if alwaysSplit:
        print "Using vtkImageDataStreamer with %d divisions"%numberOfDivisions
        streamer = vtk.vtkImageDataStreamer()
        streamer.SetNumberOfStreamDivisions(numberOfDivisions)
    elif (memLimit and not noLimits):
        print "Using vtkMemoryLimitImageDataStreamer with with limit=%dMB"%memLimit
        streamer = vtk.vtkMemoryLimitImageDataStreamer()
        streamer.SetMemoryLimit(1024*memLimit)  # per kb
    
    streamer.GetExtentTranslator().SetSplitModeToZSlab()
    
    executing = 1   # status variable
    
    streamer.SetInputConnection(pipeline.GetOutputPort())
    retval = streamer.GetOutput()
    
    if updateExtent and not wantWholeDataset:
        print "Setting update extent to ", updateExtent
        retval.SetUpdateExtent(updateExtent)
    
    print "Update extent now=",retval.GetUpdateExtent()
    streamer.Update()
    print "Done"
    executing = 0
    return retval
        

def getColorTransferFunction(color):
    """
    Convert a color to a ctf and pass a ctf through
    """
    if isinstance(color, vtk.vtkColorTransferFunction):
        return color
    ctf = vtk.vtkColorTransferFunction()
    red, green, blue = (0, 0, 0)
    ctf.AddRGBPoint(0.0, red, green, blue)
    
    red, green, blue = color
    red /= 255
    green /= 255
    blue /= 255
    ctf.AddRGBPoint(255.0, red, green, blue)    
    return ctf

def getMIP(imageData, color):
    """
    A function that will take a volume and do a simple
                 maximum intensity projection that will be converted to a
                 wxBitmap
    """   
    maxval = imageData.GetScalarRange()[1]
    imageData.SetUpdateExtent(imageData.GetWholeExtent())         

    if maxval > 255:
        shiftscale = vtk.vtkImageShiftScale()
        shiftscale.SetInputConnection(imageData.GetProducerPort())
        shiftscale.SetScale(255.0 / maxval)
        shiftscale.SetOutputScalarTypeToUnsignedChar()      
        imageData = shiftscale.GetOutput()
    
    mip = vtkJolly.vtkImageSimpleMIP()
    mip.SetInputConnection(imageData.GetProducerPort())
    
    if color == None:
        output = execute_limited(mip)
        return output

    if mip.GetOutput().GetNumberOfScalarComponents() == 1:
        ctf = getColorTransferFunction(color)
    
        maptocolor = vtk.vtkImageMapToColors()
        maptocolor.SetInputConnection(mip.GetOutputPort())
        maptocolor.SetLookupTable(ctf)
        maptocolor.SetOutputFormatToRGB()
        imagedata = execute_limited(maptocolor)
        
    else:
        imagedata = output = execute_limited(mip)
    return imagedata

def vtkImageDataToWxImage(data, sliceNumber = -1, startpos = None, endpos = None):
    """
    Converts vtk-ImageData(BMP) to a WxImage, if vtk-ImageData(DICOM), this function
    will be invalid
    """
    
    exporter = vtk.vtkImageExport()
    data.SetUpdateExtent(data.GetWholeExtent())
    data.Update()
    exporter.SetInputConnection(data.GetProducerPort())
    dataMemorySize = exporter.GetDataMemorySize()
    formatString = "%ds" % dataMemorySize
    structString = struct.pack(formatString, "")
    exporter.SetExportVoidPointer(structString)
    exporter.Export()
    width, height = data.GetDimensions()[0:2]
    image = wx.EmptyImage(width, height)
    image.SetData(structString)
    return image

if __name__ == "__main__":
    import sys
    import wx
    from jolly.datamanagement.vtkPythonMetaImageData  import *
    from jolly.visual.vtkPythonViewImage2D import *
    from vtk.util.misc import vtkGetDataRoot
    
    app = wx.PySimpleApp()
    
    image = vtkPythonMetaImageData()
    image.ReadSeriesFile("C:/head/")
    imagedata = image.getDataSet()
    flip = vtk.vtkImageFlip()
    
    if imagedata.GetNumberOfScalarComponents()>1:
        il = vtk.vtkImageLuminance()
        il.SetInput(imagedata)
        flip.SetInput(il.GetOutput())
        il.Update()
    else:
        flip.SetInput(imagedata)
    flip.SetFilteredAxis(2)
    flip.Update()
    image.setDataSet(flip.GetOutput())
    
    imagedata = getMIP(image.GetImageData(), None)
    
    view1 = vtkPythonViewImage2D()
    iren1 = vtk.vtkRenderWindowInteractor()
    view1.SetupInteractor(iren1)
    view1.SetBackground(0.0, 0.0, 0.0)
    view1.SetSliceOrientationToXY()
    view1.SetInteractorStyleTypeToNavigation()
    view1.SetInput(imagedata)
    view1.Update()
    view1.Reset()    
    view1.Render()
    iren1.Start()

    
    
    